from flask import request, jsonify, current_app
from app.api import bp
from app.services.document_service import DocumentService
from app.utils.file_handler import process_file
from app.models.knowledge_base import Document, KnowledgeBase
from app import db, logger
import os
from werkzeug.utils import secure_filename
from app.utils.validators import DocumentSplitSchema
from app.utils.errors import ValidationError, NotFoundError, ServiceError
from app.utils.response import APIResponse
from app.services.embedding_service import EmbeddingService
from app.core.vector_store import VectorStore
from app.utils.file_processors import FileProcessor
from app.utils.text_splitter import split_text
import pandas as pd
from app.utils.task_manager import task_manager
import logging

logger = logging.getLogger(__name__)

@bp.route('/document/upload/<int:kb_id>', methods=['POST'])
def upload_document(kb_id):
    """上传文档"""
    try:
        if 'file' not in request.files:
            return APIResponse.error('No file part')
            
        file = request.files['file']
        if not file:
            return APIResponse.error('No selected file')
            
        # 获取分割方法
        split_method = request.form.get('split_method', 'char')
        
        # 安全的文件名
        filename = secure_filename(file.filename)
        
        # 创建临时文件目录
        temp_dir = current_app.config['UPLOAD_FOLDER']
        os.makedirs(temp_dir, exist_ok=True)
        
        # 保存文件
        file_path = os.path.join(temp_dir, filename)
        file.save(file_path)
        
        # 创建异步任务
        task_id = task_manager.create_task(
            filename=f"处理文件: {filename}",
            kb_id=kb_id
        )
        
        # 启动异步处理
        DocumentService.process_document_async(
            file_path=file_path,
            knowledge_base_id=kb_id,
            method=split_method,
            task_id=task_id
        )
        
        return APIResponse.success({
            'task_id': task_id,
            'message': '文件上传成功，正在处理中'
        })
        
    except Exception as e:
        logger.error(f"Error uploading document: {str(e)}", exc_info=True)
        return APIResponse.error(f"上传失败: {str(e)}")

@bp.route('/document/split', methods=['POST'])
def split_document():
    data = request.get_json()
    content = data['content']
    method = data.get('method', 'char')
    params = data.get('params', {})
    
    chunks = DocumentService.split_document(content, method, **params)
    return jsonify({'chunks': chunks})

@bp.route('/document/preview', methods=['POST'])
def preview_document_split():
    """预览文档分段"""
    try:
        if 'file' not in request.files:
            raise ValidationError('No file uploaded')
        
        file = request.files['file']
        if file.filename == '':
            raise ValidationError('No selected file')
        
        # 检查文件大小
        if len(file.read()) > 10 * 1024 * 1024:  # 10MB
            raise ValidationError('File size exceeds 10MB limit')
        file.seek(0)  # 重置文件指针
        
        # 获取分割方法
        method = request.form.get('method', 'char')
        valid_methods = ['char', 'word', 'sentence', 'semantic', 'qa', 'raw']
        if method not in valid_methods:
            raise ValidationError(f'Invalid split method. Supported methods: {", ".join(valid_methods)}')
        
        # 获取文件扩展名
        file_ext = os.path.splitext(file.filename)[1].lower()
        
        # 如果是结构化文件且选择了保持原格式
        if method == 'raw' and file_ext in ['.csv', '.xlsx', '.xls']:
            try:
                if file_ext == '.csv':
                    # 保存临时文件
                    temp_path = os.path.join(current_app.config['UPLOAD_FOLDER'], 'temp_' + secure_filename(file.filename))
                    file.save(temp_path)
                    
                    try:
                        # 使用 DocumentService 的方法读取文件
                        df = DocumentService._read_structured_file(temp_path)
                    finally:
                        # 清理临时文件
                        if os.path.exists(temp_path):
                            os.remove(temp_path)
                else:
                    df = pd.read_excel(file)
                
                # 预览前5行数据
                preview_data = {
                    'total_chunks': len(df),
                    'chunks': df.head().to_dict('records'),
                    'columns': list(df.columns),
                    'format': 'structured'
                }
                return APIResponse.success(data=preview_data)
            except Exception as e:
                current_app.logger.error(f"Error previewing file: {str(e)}")
                raise ValidationError(f'Error reading file: {str(e)}')
        
        # 处理文本文件
        try:
            # 尝试不同的编码读取文本
            content = None
            encodings = ['utf-8', 'gbk', 'gb2312', 'gb18030', 'cp936', 'utf-8-sig']
            last_error = None
            
            for encoding in encodings:
                try:
                    file.seek(0)  # 重置文件指针
                    content = file.read().decode(encoding)
                    break
                except UnicodeDecodeError as e:
                    last_error = e
                    continue
            
            if content is None:
                raise ValidationError(f'Unable to decode file with supported encodings. Last error: {str(last_error)}')
            
            # 分割文本
            chunks = DocumentService.split_document(content, method)
            
            preview_data = {
                'total_chunks': len(chunks),
                'chunks': chunks[:5],  # 只返回前5个片段作为预览
                'format': 'text'
            }
            return APIResponse.success(data=preview_data)
            
        except Exception as e:
            current_app.logger.error(f"Error processing text file: {str(e)}")
            raise ValidationError(str(e))
            
    except Exception as e:
        current_app.logger.error(f"Preview split error: {str(e)}")
        if isinstance(e, ValidationError):
            raise
        raise ValidationError(str(e))

@bp.route('/document/search/<int:kb_id>', methods=['POST'])
def search_documents(kb_id):
    """搜索文档"""
    data = request.get_json()
    query = data.get('query')
    if not query:
        raise ValidationError({'query': 'Query is required'})
    
    # 生成查询向量
    query_vector = EmbeddingService.generate_embedding(query)
    
    # 搜索文档
    docs = VectorStore.search(
        query_vector,
        kb_id,
        similarity_threshold=current_app.config['VECTOR_SEARCH']['similarity_threshold'],
        max_results=current_app.config['VECTOR_SEARCH']['top_k_results']
    )
    
    return APIResponse.success(
        data=[{
            'id': doc.id,
            'content': doc.content,
            'similarity': doc.similarity
        } for doc in docs]
    ) 

@bp.route('/document/template', methods=['GET'])
def get_document_template():
    """获取文档模板格式说明"""
    templates = {
        'csv': {
            'format': 'CSV文件格式要求：\n'
                     '1. 必须包含 index 和 content 两列\n'
                     '2. content 列存放需要处理的文本内容\n'
                     '3. index 列可以是任意唯一标识',
            'example': 'index,content\n'
                      '1,"这是第一段内容"\n'
                      '2,"这是第二段内容"'
        },
        'excel': {
            'format': 'Excel文件格式要求：\n'
                     '1. 必须包含 index 和 content 两列\n'
                     '2. content 列存放需要处理的文本内容\n'
                     '3. index 列可以是任意唯一标识',
            'example': '同CSV格式，但存储为.xlsx文件'
        }
    }
    
    return APIResponse.success(
        data=templates,
        message="Document templates retrieved successfully"
    ) 

@bp.route('/preview', methods=['POST'])
def preview_split():
    try:
        if 'file' not in request.files:
            return jsonify({
                'status': 'error',
                'message': '未上传文件'
            })
            
        file = request.files['file']
        method = request.form.get('method', 'char')
        
        # 获取文件扩展名
        file_ext = os.path.splitext(file.filename)[1].lower()
        
        # 如果是结构化文件且选择了保持原格式
        if method == 'raw' and file_ext in ['.csv', '.xlsx', '.xls']:
            # 读取文件内容
            if file_ext == '.csv':
                df = pd.read_csv(file)
            else:
                df = pd.read_excel(file)
            
            # 预览前5行数据
            preview_data = df.head().to_dict('records')
            return jsonify({
                'status': 'success',
                'data': {
                    'total_chunks': len(df),
                    'chunks': preview_data,
                    'columns': list(df.columns),
                    'format': 'structured'
                }
            })
        
        # 读取文件内容
        content = file.read().decode('utf-8')
        
        # 根据方法分割文档
        chunks = DocumentService.split_document(content, method)
        
        # 对于问答对分割，需要特殊处理预览格式
        if method == 'qa':
            preview_chunks = []
            for i, chunk in enumerate(chunks, 1):
                preview_chunks.append(f"问答对 {i}:\n{chunk}")
            chunks = preview_chunks
        
        return jsonify({
            'status': 'success',
            'data': {
                'total_chunks': len(chunks),
                'chunks': chunks[:3],  # 只返回前3个片段作为预览
                'format': 'text'
            }
        })
        
    except Exception as e:
        current_app.logger.error(f"Preview split error: {str(e)}")
        return jsonify({
            'status': 'error',
            'message': f'预览失败: {str(e)}'
        }) 

@bp.route('/document/tasks', methods=['GET'])
def get_document_tasks():
    """获取文档处理任务列表"""
    tasks = task_manager.get_all_tasks()
    return APIResponse.success(
        data=[{
            'id': task.id,
            'filename': task.filename,
            'status': task.status,
            'progress': task.progress,
            'message': task.message,
            'created_at': task.created_at.isoformat(),
            'started_at': task.started_at.isoformat() if task.started_at else None,
            'completed_at': task.completed_at.isoformat() if task.completed_at else None,
            'error': task.error
        } for task in tasks]
    ) 